#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# pyre-unsafe

import json
from copy import deepcopy
from typing import Any

from augly import utils
from augly.text import intensity as txtintensity


def get_func_kwargs(
    metadata: list[dict[str, Any]] | None, local_kwargs: dict[str, Any], **kwargs
) -> dict[str, Any]:
    if metadata is None:
        return {}
    func_kwargs = deepcopy(local_kwargs)
    func_kwargs.pop("metadata")
    func_kwargs.update(**kwargs)
    return func_kwargs


def get_metadata(
    metadata: list[dict[str, Any]] | None,
    function_name: str,
    texts: list[str] | None = None,
    aug_texts: list[str] | str | None = None,
    **kwargs,
) -> None:
    if metadata is None:
        return

    assert isinstance(
        metadata, list
    ), "Expected `metadata` to be set to None or of type list"
    assert (
        texts is not None
    ), "Expected `texts` to be passed in if metadata was provided"
    assert (
        aug_texts is not None
    ), "Expected `aug_texts` to be passed in if metadata was provided"

    metadata.append(
        {
            "name": function_name,
            "input_type": "list" if isinstance(texts, list) else "string",
            "src_length": len(texts) if isinstance(texts, list) else 1,
            "dst_length": len(aug_texts) if isinstance(aug_texts, list) else 1,
            **kwargs,
        }
    )

    intensity_kwargs = {"metadata": metadata[-1], "texts": texts, **kwargs}
    metadata[-1]["intensity"] = getattr(
        txtintensity, f"{function_name}_intensity", lambda **_: 0.0
    )(**intensity_kwargs)


def get_gendered_words_mapping(mapping: str | dict[str, str]) -> dict[str, str]:
    """
    Note: The `swap_gendered_words` augmentation, including this logic, was originally
    written by Adina Williams and has been used in influential work, e.g.
    https://arxiv.org/pdf/2005.00614.pdf
    """
    assert isinstance(
        mapping, (str, dict)
    ), "Mapping must be either a dict or filepath to a mapping of gendered words"

    if isinstance(mapping, dict):
        return mapping

    if isinstance(mapping, str):
        with utils.pathmgr.open(mapping) as f:
            return json.load(f)
